from starkware.cairo.common.alloc import alloc
from starkware.cairo.common.memcpy import memcpy
from starkware.cairo.stark_verifier.air.layout import AirWithLayout, Layout, eval_oods_polynomial
from starkware.cairo.stark_verifier.air.public_input import PublicInput
from starkware.cairo.stark_verifier.air.traces import (
    CONSTRAINT_DEGREE,
    N_COMPOSITION_COLUMNS,
    TracesDecommitment,
)
from starkware.cairo.stark_verifier.core.air_interface import OodsEvaluationInfo
from starkware.cairo.stark_verifier.core.table_commitment import TableDecommitment

// Evaluates the OODS polynomial at a list of points, given a decommitment of the traces, and the
// composition commitment.
func eval_oods_boundary_poly_at_points{range_check_ptr}(
    air: AirWithLayout*,
    public_input: PublicInput*,
    eval_info: OodsEvaluationInfo*,
    n_points: felt,
    points: felt*,
    decommitment: TracesDecommitment*,
    composition_decommitment: TableDecommitment*,
) -> (evaluations: felt*) {
    alloc_locals;
    local n_original_columns = air.layout.n_original_columns;
    assert decommitment.original.n_values = n_points * n_original_columns;
    local n_interaction_columns = air.layout.n_interaction_columns;
    assert decommitment.interaction.n_values = n_points * n_interaction_columns;
    assert composition_decommitment.n_values = n_points * N_COMPOSITION_COLUMNS;
    let (evaluations: felt*) = alloc();
    eval_oods_boundary_poly_at_points_inner(
        layout=&air.layout,
        eval_info=eval_info,
        n_points=n_points,
        points=points,
        n_original_columns=n_original_columns,
        original_values=decommitment.original.values,
        interaction_values=decommitment.interaction.values,
        composition_values=composition_decommitment.values,
        evaluations=evaluations,
    );
    return (evaluations=evaluations);
}

func eval_oods_boundary_poly_at_points_inner{range_check_ptr}(
    layout: Layout*,
    eval_info: OodsEvaluationInfo*,
    n_points: felt,
    points: felt*,
    n_original_columns: felt,
    original_values: felt*,
    interaction_values: felt*,
    composition_values: felt*,
    evaluations: felt*,
) {
    alloc_locals;
    if (n_points == 0) {
        return ();
    }

    // Concatenate the values from the 3 traces into a single array of column_values.
    local n_interaction_columns = layout.n_interaction_columns;
    let (column_values) = alloc();
    memcpy(column_values, original_values, n_original_columns);
    memcpy(column_values + n_original_columns, interaction_values, n_interaction_columns);
    memcpy(
        column_values + n_original_columns + n_interaction_columns,
        composition_values,
        N_COMPOSITION_COLUMNS,
    );

    // Call autogenerated evaluation function.
    let (res) = eval_oods_polynomial(
        layout=layout,
        column_values=column_values,
        oods_values=eval_info.oods_values,
        constraint_coefficients=eval_info.constraint_coefficients,
        point=points[0],
        oods_point=eval_info.oods_point,
        trace_generator=eval_info.trace_generator,
    );
    assert evaluations[0] = res;

    return eval_oods_boundary_poly_at_points_inner(
        layout=layout,
        eval_info=eval_info,
        n_points=n_points - 1,
        points=&points[1],
        n_original_columns=n_original_columns,
        original_values=&original_values[n_original_columns],
        interaction_values=&interaction_values[n_interaction_columns],
        composition_values=&composition_values[N_COMPOSITION_COLUMNS],
        evaluations=&evaluations[1],
    );
}
